# asr.py
import asyncio
import logging
import os
from typing import Callable, Optional

import dashscope
from dashscope.audio.asr import Recognition, RecognitionCallback

logger = logging.getLogger(__name__)

# 全局设置 API Key
def init_dashscope():
    if not dashscope.api_key:
        # 你可以从环境变量读取
        dashscope.api_key = 'sk-ef1242073e404b118dfbed0987a11e51'


init_dashscope()

class ASRSession:
    """
    单个实时语音识别会话
    支持流式喂数据、流式输出、主动停止
    """




    def __init__(
        self,
        on_result: Callable[[str], None],        # 回调：返回流式文本
        on_final: Callable[[str], None],         # 回调：返回最终完整文本
        on_error: Callable[[str], None],         # 错误回调
        on_close: Callable[[], None] = None,      # 关闭回调
    ):
        self.on_result = on_result
        self.on_final = on_final
        self.on_error = on_error
        self.on_close = on_close

        self.callback = None
        self.recognition: Optional[Recognition] = None
        self.running = False
        self.task: Optional[asyncio.Task] = None
        self.loop = None  # 保存主事件循环引用

    def set_event_loop(self, loop):
        """设置主事件循环引用"""
        self.loop = loop

    def start(self):
        """启动 ASR 识别任务"""
        if self.running:
            return

        # 定义回调类
        class RealtimeASRCallback(RecognitionCallback):

            def __init__(self, asr_session):
                super().__init__()  # 调用父类构造函数
                self.asr_session = asr_session

            def on_open(self) -> None:
                logger.info("ASR 连接已打开")

            def on_close(self) -> None:
                logger.info("ASR 连接已关闭")
                if self.on_close:
                    self.asr_session.on_close()

            def on_complete(self) -> None:
                logger.info("ASR 识别完成")
                # 注意：最终文本应在 on_event 中通过 is_sentence_end 获取

            def on_error(self, message) -> None:
                error_msg = f"ASR 错误: {message.message} (request_id={message.request_id})"
                logger.error(error_msg)
                self.asr_session.on_error(error_msg)

            def on_event(self, result) -> None:
                sentence = result.get_sentence()
                print( sentence)
                if not sentence or 'text' not in sentence:
                    return

                text = sentence['text']
                self.asr_session.on_result(text)  # 流式返回

                # 如果是句子结束，也作为最终结果返回一次
                if result.is_sentence_end(sentence):
                    self.asr_session.on_final(text)

        self.callback = RealtimeASRCallback(self)
        self.recognition = Recognition(
            model='paraformer-realtime-v2',
            format='pcm',  # 假设输入是 PCM 16kHz 单声道
            sample_rate=16000,
            callback=self.callback,
            punctuation_prediction_enabled=False,
            max_sentence_silence=200,
            language_hints=['zh']
        )

        self.running = True
        self.recognition.start()

    def send_audio(self, audio_data: bytes):
        """接收一段音频数据并发送给 ASR"""
        if self.running and self.recognition:
            try:
                self.recognition.send_audio_frame(audio_data)
            except Exception as e:
                logger.error(f"发送音频帧失败: {e}")

    def stop(self):
        """主动停止 ASR 识别"""
        if self.running and self.recognition:
            try:
                self.recognition.stop()
                logger.info("ASR 识别已停止")
            except Exception as e:
                logger.error(f"停止 ASR 失败: {e}")
            finally:
                self.running = False

    def close(self):
        """关闭资源"""
        self.stop()
        self.running = False